# -*- coding: utf-8 -*-
"""
这是使用k-最近邻(KNN)算法进行人脸识别的示例。
当有大量已知面部图片时，该方式较为适用，可并行计算，进行未知人脸的预测。
"""

import math
from sklearn import neighbors
import os
import os.path
import pickle
from PIL import Image, ImageDraw
import face_recognition
from face_recognition.face_recognition_cli import image_files_in_folder

ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}


def train(train_dir, model_save_path=None, n_neighbors=None, knn_algo='ball_tree', verbose=False):
    """
    训练人脸识别分类器
    :param train_dir: 训练数据集
     需满足结构如下:
        <train_dir>/
        ├── <person1>/
        │   ├── <somename1>.jpeg
        │   ├── <somename2>.jpeg
        │   ├── ...
        ├── <person2>/
        │   ├── <somename1>.jpeg
        │   └── <somename2>.jpeg
        └── ...
    :param model_save_path: 保存分类器模型路径
    :param n_neighbors: 权值 （未指定则自动生成）
    :param knn_algo: （可选的）支持knn的数据结构. 默认 ball_tree
    :param verbose: 训练冗长
    :return: 训练好的分类器
    """
    x = []
    y = []

    # Loop through each person in the training set
    for class_dir in os.listdir(train_dir):
        if not os.path.isdir(os.path.join(train_dir, class_dir)):
            continue

        # Loop through each training image for the current person
        for img_path in image_files_in_folder(os.path.join(train_dir, class_dir)):
            image = face_recognition.load_image_file(img_path)
            face_bounding_boxes = face_recognition.face_locations(image)

            if len(face_bounding_boxes) != 1:
                # If there are no people (or too many people) in a training image, skip the image.
                if verbose:
                    print("Image {} not suitable for training: {}".format(img_path, "Didn't find a face" if len(
                        face_bounding_boxes) < 1 else "Found more than one face"))
            else:
                # Add face encoding for current image to the training set
                x.append(face_recognition.face_encodings(image, known_face_locations=face_bounding_boxes)[0])
                y.append(class_dir)

    # Determine how many neighbors to use for weighting in the KNN classifier
    if n_neighbors is None:
        n_neighbors = int(round(math.sqrt(len(x))))
        if verbose:
            print("Chose n_neighbors automatically:", n_neighbors)

    # Create and train the KNN classifier
    knn_clf = neighbors.KNeighborsClassifier(n_neighbors=n_neighbors, algorithm=knn_algo, weights='distance')
    knn_clf.fit(x, y)

    # Save the trained KNN classifier
    if model_save_path is not None:
        with open(model_save_path, 'wb') as f:
            pickle.dump(knn_clf, f)

    return knn_clf


def predict(x_img_path, knn_clf=None, model_path=None, distance_threshold=0.6):
    """
    使用knn分类器识别图片中的人脸
    :param x_img_path: 图片路径
    :param knn_clf: (optional) a knn classifier object. if not specified, model_save_path must be specified.
    :param model_path: (optional) path to a pickled knn classifier. if not specified, model_save_path must be knn_clf.
    :param distance_threshold: 识别阈值
    :return: 所有的人名和对应的人脸位置列表: [(name, bounding box), ...].未识别出的为 'unknown'
    """
    if not os.path.isfile(x_img_path) or os.path.splitext(x_img_path)[1][1:] not in ALLOWED_EXTENSIONS:
        raise Exception("Invalid image path: {}".format(x_img_path))

    if knn_clf is None and model_path is None:
        raise Exception("Must supply knn classifier either thourgh knn_clf or model_path")

    # Load a trained KNN model (if one was passed in)
    if knn_clf is None:
        with open(model_path, 'rb') as f:
            knn_clf = pickle.load(f)

    # Load image file and find face locations
    x_img = face_recognition.load_image_file(x_img_path)
    x_face_locations = face_recognition.face_locations(x_img)

    # If no faces are found in the image, return an empty result.
    if len(x_face_locations) == 0:
        return []

    # Find encodings for faces in the test iamge
    faces_encodings = face_recognition.face_encodings(x_img, known_face_locations=x_face_locations)

    # Use the KNN model to find the best matches for the test face
    closest_distances = knn_clf.kneighbors(faces_encodings, n_neighbors=1)
    are_matches = [closest_distances[0][i][0] <= distance_threshold for i in range(len(x_face_locations))]

    # Predict classes and remove classifications that aren't within the threshold
    return [(pred, loc) if rec else ("unknown", loc) for pred, loc, rec in
            zip(knn_clf.predict(faces_encodings), x_face_locations, are_matches)]


def show_prediction_labels_on_image(img_path, predictions):
    """
    可视化显示分类结果
    :param img_path: 图片路径
    :param predictions: 分类器分类结果
    :return:
    """
    pil_image = Image.open(img_path).convert("RGB")
    draw = ImageDraw.Draw(pil_image)

    for name, (top, right, bottom, left) in predictions:
        # Draw a box around the face using the Pillow module
        draw.rectangle(((left, top), (right, bottom)), outline=(0, 0, 255))

        # There's a bug in Pillow where it blows up with non-UTF-8 text
        # when using the default bitmap font
        name = name.encode("UTF-8")

        # Draw a label with a name below the face
        text_width, text_height = draw.textsize(name)
        draw.rectangle(((left, bottom - text_height - 10), (right, bottom)), fill=(0, 0, 255), outline=(0, 0, 255))
        draw.text((left + 6, bottom - text_height - 5), name, fill=(255, 255, 255, 255))

    # Remove the drawing library from memory as per the Pillow docs
    del draw

    # Display the resulting image
    pil_image.show()


def re_trained(train_path="knn_examples/train", model_save_path="trained_knn_model.clf"):
    print("开始训练 KNN classifier...")
    train(train_path, model_save_path, n_neighbors=2)
    print("Training complete!")


if __name__ == "__main__":
    # knn模型路径 训练数据路径
    knn_model_path = "trained_knn_model.clf"
    knn_train_path = "knn_examples/train"

    # 未知人脸的图片
    full_file_path = "knn_examples/test/wyf_1.jpg"

    # STEP 1:训练KNN分类器 并保存到本地  不需要每次都执行
    # re_trained(train_path=knn_train_path, model_save_path=knn_model_path)

    # 使用经过训练的分类器模型查找图像中的所有人
    predictions_res = predict(full_file_path, model_path=knn_model_path, distance_threshold=0.49)

    # 输出结果
    # for name, (top, right, bottom, left) in predictions_res:
    #     print("- Found {} at ({}, {})".format(name, left, top))

    # 显示结果（可视化）
    show_prediction_labels_on_image(full_file_path, predictions_res)
